Fix device selection on Apple Silicon (MPS) - #205

Summary

On Apple Silicon (Mac, MPS), several nodes selected their device with a module-level

device = "cuda" if torch.cuda.is_available() else "cpu"

Since torch.cuda.is_available() is False on a Mac, they silently fell back to CPU even when ComfyUI itself had selected MPS. The model still runs, just far slower, so there's no error — only a job that sits at 0% for minutes, which reads as a hang.

This defers device selection to ComfyUI's own device management via a small shared helper, so the nodes follow whatever backend ComfyUI chose (CUDA, MPS, or CPU) — the same approach the SAM2 / SAM3 / Florence2 / Segment / SDMatte nodes already use.

Changes

• Add get_device() to py/AILab_utils.py, wrapping comfy.model_management.get_torch_device().
• Use it in the nodes that hard-coded the CUDA-or-CPU check: BiRefNet, RMBG, ClothSegment, BodySegment, FashionSegment, FaceSegment, SegmentV2.
• Centralizing it in one helper keeps the check from drifting again.
• LamaRemover already uses get_torch_device() and is left unchanged.

device is only ever used as .to(device) / map_location=device / a hashable cache key, so returning a torch.device object instead of a string is safe at every call site.

Evidence (before/after, Apple Silicon)

• Before (CPU fallback): a single BiRefNet image ≈ 18 min.
• After (MPS): a 21-second video processed in ≈ 9 min (~1s/frame).

Environment: Apple Silicon, macOS 15.7.4, 48 GB unified memory, ComfyUI 0.20.1 (Desktop), PyTorch 2.10.0, Python 3.13, ComfyUI-RMBG v3.0.0. ComfyUI startup correctly reports Device: mps.

Notes

• A few ops may not yet be implemented for MPS. If you hit one, launch ComfyUI with PYTORCH_ENABLE_MPS_FALLBACK=1 so those fall back to CPU instead of erroring.

Related issues

• Fixes BiRefNet node stuck at 0% with no error on Mac Apple Silicon (ComfyUI Desktop) #200 — "BiRefNet node stuck at 0% with no error on Mac Apple Silicon (ComfyUI Desktop)". This is exactly the reported symptom (BiRefNet stuck at 0% while RMBG-2.0 works on the same machine). The issue was auto-closed by the stale bot after 30 days of inactivity, not by a fix.
• Related to Suggestion: let user choose to use "cuda" or "cpu" as pytorch device. #91 — "let user choose 'cuda' or 'cpu' as pytorch device", which first flagged this exact device = "cuda" if torch.cuda.is_available() else "cpu" pattern. A device dropdown was later added to the SAM2 node, but these image-segmentation nodes (BiRefNet, RMBG, Cloth/Body/Fashion/FaceSegment, SegmentV2) still hard-coded the CUDA-or-CPU check until now.
• Related to SAM3 Segmentation - CPU/GPU device selection (issue) / Model unloading (feature request) #135 — SAM3 CPU/GPU device-selection crash; another data point that device handling has been a weak spot. (SAM3 already uses get_torch_device() and is unchanged here. Also closed for inactivity.)

If you'd prefer a user-facing device dropdown (Auto/CUDA/CPU/MPS) on these nodes — matching the YoloV8/SAM3 pattern — rather than this automatic deferral, happy to follow up with that.